//
//  patchcgmemory.m
//  patchcg
//
//  this doesn't work exactly as advertised, but it should be enough to get
//  you going.
//  
//  Created by thrust on 25/01/06.
//  Copyright 2006 __MyCompanyName__. All rights reserved.
//

#import "patchcgmemory.h"
#include "GetBSDProcessList.h"
#import "NSProcInfo.h"
#import "NSMachOTaskImage.h"
#import "NSMachOSymbolTable.h"
//#include "vmutils_structs.h"					// for <nlist.h>
#import "NSFunctionNameDemangler.h"

/*! ----------------------------+---------------+-------------------------------
	@discussion
	search for the windowserver pid. can we assume there is ever only one?
*/
int getWindowServerPid() {
	kinfo_proc *procList = NULL;
	size_t procCount = 0;
	int err;
	int pid = 0;
	
	NSLog(@"searching for windowserver pid");
	err = GetBSDProcessList(&procList, &procCount);
	
	int idx;
	for (idx = 0; idx < procCount; idx++) {
		//printf("process: %s\n", procList[idx].kp_proc.p_comm);
		if (strcmp("WindowServer", procList[idx].kp_proc.p_comm) == 0) {
			pid = procList[idx].kp_proc.p_pid;
			NSLog(@"windowserver found, pid: %d\n", pid);
			break;
		}
	}
	free(procList);
	
	printf("\n");
	return pid;
}

/*! ----------------------------+---------------+-------------------------------
	@discussion
	convert pid to task
*/
vm_map_t
getWindowServerTask(int pid) {
	vm_map_t task;
	kern_return_t kernErr;
	kernErr = task_for_pid(current_task(), pid, &task);
	
	if (kernErr != KERN_SUCCESS) {
		fprintf(stderr, "task_for_pid err: %d\n", kernErr);
		// do we want to bail here?
	}
	
	return task;
}

UInt32
getCGXLFindBestPositionForDisplayAddress(int pid) {
	if (pid == 0) {
		return 0;
	}
	
	NSLog(@"creating task image of windowserver");
	NSProcInfo *theProcInfo;
	theProcInfo = [[NSProcInfo alloc] initWithPid:pid];
	
	NSMachOTaskImage *theTaskImage;
	theTaskImage = [[NSMachOTaskImage alloc] initWithTask:[theProcInfo task]];
	NSLog(@"finished creating NSMachOTaskImage object");
//	
//	NSArray *theHeaderArray;
//	theHeaderArray = [theTaskImage getMachHeaders];
//	NSLog(@"there are %d mach-o headers in the task image", [theHeaderArray count]);
//	
//	NSEnumerator *theEnum;
//	theEnum = [theHeaderArray objectEnumerator];
//	
//	// find coregraphics
//	NSLog(@"searching for coregraphics");
//	NSMachOHeader *theHeader;
//	NSMachOSymbolTable *theSymbolTable;
//	BOOL cgfound = false;
//	
//	while (theHeader = [theEnum nextObject]) {
//		theSymbolTable = [theTaskImage symbolTableForHeader: theHeader];
//		
//		if ([[[theHeader name] lastPathComponent] compare:@"CoreGraphics"] == NSOrderedSame) {
//			// coregraphics found
//			cgfound = true;
//			NSLog(@"coregraphics found");
//			break;
//		}
//	}
//	if (cgfound) {
//		NSLog(@"symtable for: %@", [[theHeader name] lastPathComponent]);
//		NSLog(@"symtable contains %d symbols", [theSymbolTable numberOfSymbols]);
//		NSLog(@"symtable slide value: 0x%08x", [theHeader slideValue]);
//	} else {
//		// no cg?
//		NSLog(@"could not find coregraphics");
//		return 0;
//	}
//	
//	// find CGXLFindBestPositionForDisplay()
//	NSLog(@"searching for CGXLFindBestPositionForDisplay");
//	struct nlist *theSymArray;
//	theSymArray = [theSymbolTable symbols];
//	int symIdx;
//	char *symName;
//	unsigned long symAddr = 0;
//	
//	for (symIdx = 0; symIdx < [theSymbolTable numberOfSymbols]; symIdx++) {
//		symName = theSymArray[symIdx].n_name;
//		//NSLog(@"sym %d: 0x%08x %s", symIdx, theSymArray[symIdx].n_value, theSymArray[symIdx].n_name);
//		if (strcmp(symName, "_CGXLFindBestPositionForDisplay") == 0) {
//			symAddr = theSymArray[symIdx].n_value;
//			break;
//		}
//	}
//	if (symAddr != 0) {
//		// symbol found
//		NSLog(@"symbol %s found at 0x%08x", symName, symAddr);
//		if ([theHeader slideValue] != 0) {
//			NSLog(@"adding slide amount, untested");
//			symAddr += [theHeader slideValue];
//		}
//	} else {
//		// _CGXLFindBestPositionForDisplay not found
//		NSLog(@"_CGXLFindBestPositionForDisplay not found");
//		return 0;
//	}
	unsigned long symAddr = 0;
	
	// this may not be the most efficient way to search for a particular symbol,
	// since we have to search all of the loaded images, but it's certainly the
	// most readable, and seems plenty quick enough.
	symAddr = [theTaskImage addressOfSymbolWithCName:"CGXLFindBestPositionForDisplay"];
	NSLog(@"addr 0x%08x", symAddr);

	return symAddr;
}

/*! ----------------------------+---------------+-------------------------------
	@discussion
	check for old operation, if it's as expected then replace with new operation
*/
int
patchMemory(vm_map_t task, UInt32 addr, UInt32 oldOp, UInt32 newOp) {
	kern_return_t result;
	
	// turn on write access
	result = vm_protect(task, addr, sizeof(UInt32), false, VM_PROT_ALL | VM_PROT_COPY);
	if (result == KERN_SUCCESS) {
		// ok to write to the region
		NSLog(@"vm_protect success");
	} else {
		NSLog(@"vm_protect error: %d", result);
		result = vm_protect(task, addr, sizeof(UInt32), false, VM_PROT_DEFAULT | VM_PROT_COPY);
		
		if (result == KERN_SUCCESS) {
			NSLog(@"vm_protect success");
		} else {
			NSLog(@"vm_protect error: %d", result);
		}
	}

	//int err = ENOENT;
	
	return 0;
}

/*! ----------------------------+---------------+-------------------------------
	@discussion
*/
int
patchCGmemory() {
	NSLog(@"patchCGmemory starting");
	
	// get the windowserver pid
	int pid = getWindowServerPid();
	
	if (pid == 0) {
		NSLog(@"windowserver process not found, exiting");
		return 0;
	}
	
	// look up CGXL.. address via vmutil private framework
	UInt32 symAddr = getCGXLFindBestPositionForDisplayAddress(pid);
	
	if (symAddr == 0) {
		return 0;
	}
	
	printf("\n");
	NSLog(@"proceeding with patch");
	vm_map_t task = getWindowServerTask(pid);		
	NSLog(@"windowserver task: %d\n", task);
	
	// patch task
	vm_address_t CGXLFindBestPositionForDisplayAddr = symAddr;
	vm_address_t regionAddr = CGXLFindBestPositionForDisplayAddr;
	//unsigned attribs = 0;
	kern_return_t result;
	vm_size_t regionSize = 0;
	vm_region_basic_info_data_t vminfo;
	mach_msg_type_number_t infoCnt = VM_REGION_BASIC_INFO_COUNT;
	mach_port_t object_name = 0;
	
	result = vm_region(task, &regionAddr, &regionSize, VM_REGION_BASIC_INFO,
						(vm_region_info_t)(&vminfo), &infoCnt, &object_name);
	
	if (result == KERN_SUCCESS) {
		// get the attributes
		if ( vminfo.protection & VM_PROT_READ ) {
			//attribs |= VMREGION_READABLE;
			NSLog(@"vm readable");
		}
		if ( vminfo.protection & VM_PROT_WRITE ) {
			//attribs |= VMREGION_WRITABLE;
			NSLog(@"vm writable");
		}
		if ( vminfo.protection & VM_PROT_EXECUTE ) {
			//attribs |= VMREGION_EXECUTABLE;
			NSLog(@"vm exec");
		}

	}
	
	// calculate the offset into the region of our symbol
	vm_address_t CGXLFindBestPositionForDisplayOffset = CGXLFindBestPositionForDisplayAddr - regionAddr;
	
	NSLog(@"CGXLFindBestPositionForDisplayOffset: %08x", CGXLFindBestPositionForDisplayOffset);
	
	char *buffer;
	// read a whole region worth?
	vm_size_t bufsize = CGXLFindBestPositionForDisplayOffset + 32;
	vm_size_t actualSize;
	buffer = (void *) malloc(bufsize);
	if (!buffer) {
		// error
		NSLog(@"could not create buffer of size %d", bufsize);
	}
	
	// grab the region into our local buffer
	result = vm_read_overwrite(task, regionAddr, bufsize, (vm_address_t) buffer, &actualSize);
	
	if (result == KERN_SUCCESS) {
		NSLog(@"vm_read_overwrite success");
	} else {
		NSLog(@"vm_read_overwrite error: %d", result);
	}
	
	// dump the memory at our symbol location
	int idx;
	for (idx = 0; idx < 4; idx++) {
		UInt32 *op;
		op = (UInt32 *)(buffer + CGXLFindBestPositionForDisplayOffset + idx * sizeof(UInt32));
		NSLog(@"op: %08x", *op);
	}
	// dump the buffer
	free(buffer);
	
	// turn on write access
	result = vm_protect(task, regionAddr, regionSize, false, VM_PROT_ALL | VM_PROT_COPY);
	if (result == KERN_SUCCESS) {
		// ok to write to the region
		NSLog(@"vm_protect success");
	} else {
		NSLog(@"vm_protect error: %d", result);
		result = vm_protect(task, regionAddr, regionSize, false, VM_PROT_DEFAULT | VM_PROT_COPY);
		
		if (result == KERN_SUCCESS) {
			NSLog(@"vm_protect success");
		} else {
			NSLog(@"vm_protect error: %d", result);
		}
	}
	
	// at this point, we can write to the region?
//	UInt32 op1 = 0x38600000;						// li r3, 0x0
//	UInt32 op2 = 0x4e800020;						// blr
//	
//	result = vm_write(task, CGXLFindBestPositionForDisplayAddr, (vm_address_t) &op1, sizeof(op1));
//	if (result == KERN_SUCCESS) {
//		NSLog(@"vm_write op1 success");
//	} else {
//		NSLog(@"vm_write op1 error: %d", result);
//	}
//	result = vm_write(task, CGXLFindBestPositionForDisplayAddr + 4, (vm_address_t) &op2, sizeof(op1));
//	if (result == KERN_SUCCESS) {
//		NSLog(@"vm_write op2 success");
//	} else {
//		NSLog(@"vm_write op2 error: %d", result);
//	}

	// verify the changes
	buffer = (void *) malloc(bufsize);
	if (!buffer) {
		// error
		NSLog(@"could not create buffer of size %d", bufsize);
	}
	
	// grab the region into our local buffer
	result = vm_read_overwrite(task, regionAddr, bufsize, (vm_address_t) buffer, &actualSize);
	
	if (result == KERN_SUCCESS) {
		NSLog(@"vm_read_overwrite success");
	} else {
		NSLog(@"vm_read_overwrite error: %d", result);
	}
	
	// dump the memory at our symbol location
	for (idx = 0; idx < 4; idx++) {
		UInt32 *op;
		op = (UInt32 *)(buffer + CGXLFindBestPositionForDisplayOffset + idx * sizeof(UInt32));
		NSLog(@"op: %08x", *op);
	}
	// dump the buffer
	free(buffer);

	return 0;
}


